import os
import time
import hashlib
import datetime
import pandas as pd
from PIL import Image
import pytesseract
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from openpyxl import load_workbook

# ChromeDriver 设置
chrome_driver_path = "/path/to/your/chromedriver"#设置为你的chromedriver路径
service = Service(chrome_driver_path)
options = webdriver.ChromeOptions()
options.add_argument("--headless")
options.add_argument("--start-maximized")
driver = webdriver.Chrome(service=service, options=options)

# 读取原始 Excel 文件
file_path = "Article_Summary.xlsx"
df = pd.read_excel(file_path)

# 确保两列存在
if "Screenshot Path" not in df.columns:
    df["Screenshot Path"] = ""
if "OCR Text" not in df.columns:
    df["OCR Text"] = ""

# 创建截图文件夹
screenshot_folder = "screenshots"
os.makedirs(screenshot_folder, exist_ok=True)

# Excel 输出文件
output_excel_file = "Article_Summary_with_Screenshots.xlsx"
if os.path.exists(output_excel_file):
    writer = pd.ExcelWriter(output_excel_file, engine="openpyxl", mode="a", if_sheet_exists="overlay")
    book = load_workbook(output_excel_file)
else:
    writer = pd.ExcelWriter(output_excel_file, engine="openpyxl")
    book = None

# 主循环：截图 + OCR + 写入
for index, row in df.iterrows():
    url = row["Link"]
    try:
        driver.get(url)
        time.sleep(2)

        # 文件名：时间戳 + md5
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        url_hash = hashlib.md5(url.encode("utf-8")).hexdigest()[:8]
        screenshot_filename = f"{timestamp}_{url_hash}.png"
        screenshot_path = os.path.join(screenshot_folder, screenshot_filename)

        # 页面尺寸并截图
        width = driver.execute_script("return document.body.scrollWidth")
        height = driver.execute_script("return document.body.scrollHeight")
        driver.set_window_size(width, height)
        driver.save_screenshot(screenshot_path)

        # 相对路径写入
        screenshot_rel_path = os.path.relpath(screenshot_path)
        df.at[index, "Screenshot Path"] = screenshot_rel_path
        print(f"✅ 截图完成：{screenshot_rel_path}")

        # OCR 文字识别
        try:
            img = Image.open(screenshot_path)
            text = pytesseract.image_to_string(img, lang="chi_sim")
            df.at[index, "OCR Text"] = text.strip()
            print(f"🧠 OCR 成功：{text[:30]}..." if text else "无文字识别结果")
        except Exception as ocr_err:
            print(f"❌ OCR 失败: {ocr_err}")
            df.at[index, "OCR Text"] = "OCR Failed"

        # 写入 Excel（首行带标题）
        if index == 0 and book is None:
            df.iloc[[index]].to_excel(writer, sheet_name="Sheet1", index=False, header=True, startrow=index)
        else:
            df.iloc[[index]].to_excel(writer, sheet_name="Sheet1", index=False, header=False, startrow=index + 1)

        writer._save()

    except Exception as e:
        print(f"❌ 页面处理失败: {url}, 错误: {e}")
        df.at[index, "Screenshot Path"] = "Failed"
        df.at[index, "OCR Text"] = "Failed"
        df.iloc[[index]].to_excel(writer, sheet_name="Sheet1", index=False, header=False, startrow=index + 1)
        writer._save()

# 关闭资源
driver.quit()
writer.close()

print(f"\n🎉 所有处理完成，文件输出至：{output_excel_file}")